import random

import pymongo
from 代理.domain import Proxy
from 代理.setting import MONGO_URL
from 代理.untils.log import logger


class MONGOPOOL():
    def __init__(self):
        self.client = pymongo.MongoClient(MONGO_URL)
        self.proxies = self.client['proxies_pool']['proxies']

    def __del__(self):
        self.client.close()

    def insert_proxy(self, proxy):
        count = self.proxies.count_documents({'_id': proxy.ip})
        if count == 0:
            dic = proxy.__dict__
            dic['_id'] = proxy.ip
            self.proxies.insert_one(dic)
            logger.warning("插入新的代理：{}".format(proxy.ip))
        else:
            logger.warning("已经存在代理:{}".format(proxy.ip))

    def update_one(self, proxy):
        self.proxies.update_one({'_id': proxy.ip}, {'$set': proxy.__dict__})

    def delete_one(self, proxy):
        self.proxies.delete_one({'_id': proxy.ip})
        logger.info("删除代理ip：{}".format(proxy.ip))

    def find_all(self):
        cursor = self.proxies.find()
        for item in cursor:
            item.pop('_id')
            proxy = Proxy(**item)
            yield proxy

    def find(self, conditions={}, count=0):
        '''
        :param conditions: 查询条件字典
        :param count: 限制取出多少个代理ip
        :return: 返回proxy对象
        '''
        cursor = self.proxies.find(conditions, limit=count).sort(
            [('score', pymongo.DESCENDING), ('speed', pymongo.ASCENDING)])
        proxy_list = []
        for item in cursor:
            item.pop('_id')
            proxy = Proxy(**item)
            proxy_list.append(proxy)
        return proxy_list

    def get_proxies(self, protocal=None, domain=None, count=0, nick_type=0):
        '''
        :param protocal: 协议类型http,https
        :param domain: 域名
        :param count: 限制多少个代理Ip
        :param nick_type: 获取高匿名代理Ip
        :return:
        '''
        conditions = {'nick_type': nick_type}
        if protocal is None:
            conditions['protocol'] = 2
        elif protocal.lower() == 'http':
            conditions['protocol'] = {'$in': [1, 2]}
        else:
            conditions['protocol'] = {'$in': [0, 2]}
        if domain:
            conditions['no_domain'] = {'$nin': [domain]}
        return self.find(conditions, count=count)

    def random_proxies(self, protocal=None, domain=None, count=0, nick_type=0):
        proxy_list = self.get_proxies(protocal=protocal, domain=domain, count=count, nick_type=nick_type)
        return random.choice(proxy_list)


if __name__ == '__main__':
    mongo = MONGOPOOL()
    proxy = Proxy(ip='222.116.112.4', port='7757')
    mongo.insert_proxy(proxy)
    # for proxy in mongo.find_all():
    #     print(proxy)
    # for proxy in mongo.find({'protocol':2}): #条件查询
    #     print(proxy)
    # for proxy in mongo.get_proxies(protocal='http'):
    #     print(proxy)
